#include <algorithm>
#include "pub_defs.h"
#include <mpi.h>

void gameoflife_optimized(uint8_t* local_arr, const int m, const int n, const int rank, const int size) {
    uint8_t* tmp_arrs[2];
    tmp_arrs[0] = new uint8_t[(m+2)*(n+2)];
    tmp_arrs[1] = new uint8_t[(m+2)*(n+2)];
    std::copy(local_arr, local_arr+(m+2)*(n+2), tmp_arrs[0]);
    for (int gen=0;gen<GEN;gen++){
        int tg = gen%2;
        int ng = (gen+1)%2;
        for (int i = 1; i < m+1; i++) {
            for (int j = 1; j < n+1; j++) {
                int neightbours = 
                    VA(tmp_arrs[tg], i-1, j-1, n)
                    +VA(tmp_arrs[tg], i-1, j, n)
                    +VA(tmp_arrs[tg], i-1, j+1, n)
                    +VA(tmp_arrs[tg], i, j-1, n)
                    +VA(tmp_arrs[tg], i, j+1, n)
                    +VA(tmp_arrs[tg], i+1, j-1, n)
                    +VA(tmp_arrs[tg], i+1, j, n)
                    +VA(tmp_arrs[tg], i+1, j+1, n);
                
                switch (neightbours) {
                    case 0:
                        VA(tmp_arrs[ng], i, j, n) = 0;
                        break;
                    case 1:
                        VA(tmp_arrs[ng], i, j, n) = 0;
                        break;
                    case 2:
                        VA(tmp_arrs[ng], i, j, n) = VA(tmp_arrs[tg], i, j, n);
                        break;
                    case 3:
                        VA(tmp_arrs[ng], i, j, n) = 1;
                        break;
                    default:
                        VA(tmp_arrs[ng], i, j, n) = 0;
                }
            }    
        }
        // 发送接收数据
        // 先发送尾部数据到下一个rank， 从上一个rank接收顶部数据
        MPI_Sendrecv(&tmp_arrs[ng][(1+m-1)*(n+2)], n+2, MPI_UNSIGNED_CHAR, (rank+1)%size, 0, 
                    &tmp_arrs[ng][0], n+2, MPI_UNSIGNED_CHAR, (rank+size-1)%size, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
        
        // 发送顶部数据到上一个rank，从下一个rank接收尾部数据
        MPI_Sendrecv(&tmp_arrs[ng][1*(n+2)], n+2, MPI_UNSIGNED_CHAR, (rank+size-1)%size, 0, 
                    &tmp_arrs[ng][(1+m)*(n+2)], n+2, MPI_UNSIGNED_CHAR, (rank+1)%size, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
        // printf("\nafter all send recv\n");
        if (rank == 0){
            std::fill_n(tmp_arrs[ng], n+2, 0);
        }
        if (rank == size-1) {
            std::fill_n(tmp_arrs[ng]+((m+1)*(n+2)), n+2, 0);
        }

        // debug info
        // sleep(rank*1);
        // std::cout<< "rank: " << rank << std::endl;
        // print_arr(tmp_arrs[ng], m+2, N+2);
    }
    int endindex = (GEN)%2;
    std::copy(tmp_arrs[endindex], tmp_arrs[endindex]+(m+2)*(n+2), local_arr);
}